"""数据处理进程模块 - 异步计算统计和生成报告

数据格式说明
============

输入数据格式（从共享内存列表读取）：

1. latency_samples (延迟数据)
   - 类型: list[float]
   - 单位: 毫秒 (ms)
   - 内容: 每个元素是一次推理的延迟时间
   - 示例: [3.245, 3.512, 3.102, 3.876, 4.231, ...]
   - 来源: infer_process.py 中每次 runner.infer() 的执行时间
   - 说明: 列表长度等于 num_inputs（推理轮数）

2. power_samples (功耗采样数据)
   - 类型: list[tuple[float, float]]
   - 单位: (毫秒, 毫瓦) - (ms, mW)
   - 内容: 每个元组是 (timestamp_ms, power_mW)
     * timestamp_ms: 相对时间戳，从采样开始时刻起计算（毫秒）
     * power_mW: 瞬时功率（毫瓦）
   - 示例: [(0.0, 18500.0), (20.5, 19200.0), (41.2, 22300.0), (60.8, 21800.0), ...]
   - 来源: monitor_process.py 中 PowerSampler 的后台采样线程
   - 说明: 
     * 采样频率由 sample_rate_hz 决定（默认 50Hz，即每 20ms 一个样本）
     * 第一个样本的时间戳通常为 0.0 或接近 0.0
     * 至少需要 2 个样本才能计算能量消耗（energy_J）
     * 如果功耗采样失败或未启用，列表为空

输出数据格式（放入 result_queue）：

3. summary (统计结果)
   - 类型: Dict[str, Any]
   - 内容: 包含延迟统计、功耗统计、性能指标等
   - 示例:
     {
       "lat_avg_ms": 3.543,
       "lat_min_ms": 2.546,
       "lat_max_ms": 7.351,
       "p50_ms": 3.130,
       "p90_ms": 5.358,
       "p95_ms": 6.323,
       "p99_ms": 7.062,
       "rss_mb": 40.046,
       "init_ms": 718.130,
       "qps": 274.950,
       "batch": 1,
       "inputs": 100,
       "power_avg_W": 22.991,
       "power_peak_W": 34.396,
       "power_min_W": 18.783,
       "energy_J": 42.694,
       "num_power_samples": 93,
       "perf_per_W": 11.958
     }
"""

from __future__ import annotations

import time
from typing import Any, Dict
from multiprocessing import Queue
from multiprocessing.synchronize import Barrier as BarrierType, Event as EventType
from multiprocessing.managers import ListProxy

import numpy as np


def compute_metrics(
    latencies: list[float],
    power_data: list[tuple[float, float]],
    batch_size: int,
    num_inputs: int,
    wall_time: float,
    init_ms: float
) -> Dict[str, Any]:
    """
    计算性能统计指标
    
    Args:
        latencies: 延迟列表（毫秒）
        power_data: 功耗采样数据 [(timestamp_ms, power_mW), ...]
        batch_size: 批处理大小
        num_inputs: 推理轮数
        wall_time: 总墙钟时间（秒）
        init_ms: 初始化时间（毫秒）
        
    Returns:
        统计结果字典
    """
    import psutil
    
    # 延迟统计
    lat = np.array(latencies, dtype=np.float64)
    p = lambda q: float(np.percentile(lat, q)) if lat.size > 0 else 0.0
    
    # RSS 内存占用
    rss_mb = psutil.Process().memory_info().rss / (1024 * 1024)
    
    # 基础性能指标
    result = {
        "lat_avg_ms": float(lat.mean()) if lat.size else 0.0,
        "lat_min_ms": float(lat.min()) if lat.size else 0.0,
        "lat_max_ms": float(lat.max()) if lat.size else 0.0,
        "p50_ms": p(50),
        "p90_ms": p(90),
        "p95_ms": p(95),
        "p99_ms": p(99),
        "rss_mb": float(rss_mb),
        "init_ms": init_ms,
        "qps": (num_inputs * batch_size) / wall_time if wall_time > 0 else 0.0,
        "batch": batch_size,
        "inputs": num_inputs,
    }
    
    # 功耗统计
    # 注意：至少需要1个采样点才能计算平均功率，至少需要2个点才能计算能量
    if power_data and len(power_data) >= 1:
        power_sum_mW = 0.0
        power_max_mW = 0.0
        power_min_mW = float('inf')
        energy_J = None  # 默认 None，只有2个以上样本才能计算
        
        for i in range(len(power_data)):
            ts_ms, p_mW = power_data[i]
            power_sum_mW += p_mW
            power_max_mW = max(power_max_mW, p_mW)
            power_min_mW = min(power_min_mW, p_mW)
            
            # 只有2个以上样本才能计算能量
            if i > 0:
                dt_ms = ts_ms - power_data[i - 1][0]
                if dt_ms > 0:  # 确保时间差有效
                    p_avg_mW = (p_mW + power_data[i - 1][1]) / 2.0
                    if energy_J is None:
                        energy_J = 0.0
                    energy_J += (p_avg_mW / 1000.0) * (dt_ms / 1000.0)
        
        power_avg_W = power_sum_mW / (1000.0 * len(power_data))
        
        result.update({
            "power_avg_W": power_avg_W,
            "power_peak_W": power_max_mW / 1000.0,
            "power_min_W": power_min_mW / 1000.0,
            "energy_J": energy_J,  # 如果只有1个样本，这里是 None
            "num_power_samples": len(power_data),
        })
        
        # 能效指标（需要平均功率 > 0）
        if power_avg_W > 0:
            result["perf_per_W"] = result["qps"] / power_avg_W
    else:
        result.update({
            "power_avg_W": None,
            "power_peak_W": None,
            "power_min_W": None,
            "energy_J": None,
            "num_power_samples": 0,
        })
    
    return result


def data_handler_worker(
    barrier: BarrierType,
    event_list: list[EventType],
    latency_samples: ListProxy,
    power_samples: ListProxy,
    result_queue: Queue,
    metadata_queue: Queue
):
    """
    数据处理进程工作函数 - 持续运行，复用于所有模型测试
    
    Args:
        barrier: 三进程同步屏障
        event_list: 事件列表
            [0]: model_infer_stop_event - 推理完成信号
            [1]: data_copy_completed_event - 数据复制完成信号  
            [2]: monitor_stop_event - 全局退出信号
            [3]: power_data_ready_event - 功耗数据就绪信号（等待此信号后再读取数据）
        latency_samples: 共享列表，存储延迟数据
        power_samples: 共享列表，存储功耗采样数据
        result_queue: 结果队列，返回统计结果
        metadata_queue: 元数据队列，接收测试配置信息
    """
    while True:
        # 等待三进程同步
        barrier.wait()
        
        # 等待推理完成
        event_list[0].wait()
        
        # **关键修复**：等待监控进程完成数据复制并发出就绪信号
        # 使用信号量机制确保数据处理进程在监控进程之后运行
        # 只有当监控进程完成数据复制并设置 event_list[3] 后，才能读取数据
        if len(event_list) > 3:
            # 如果存在 event_list[3]（功耗数据就绪信号），等待它
            # 这样确保数据处理进程在监控进程完成数据复制后才读取
            # 等待时间最多5秒，避免死锁
            # 注意：只有在启用功耗采样时，event_list[3] 才会被设置
            event_list[3].wait(timeout=5.0)
        
        # 深拷贝数据（避免竞态条件）
        latencies = list(latency_samples)
        power_data = list(power_samples)
        
        # 统一的状态打印（避免与监控进程输出冲突）
        print(f"    [Data Handler] Latency samples: {len(latencies)}, Power samples: {len(power_data)}")
        if len(power_data) == 0:
            print(f"    [WARNING] No power data collected - power sampling may be disabled or failed")
        elif len(power_data) < 2:
            print(f"    [WARNING] Only {len(power_data)} power sample(s) - need at least 2 for energy calculation")
        
        # 获取测试元数据（批处理大小、输入数等）
        try:
            metadata = metadata_queue.get(timeout=1.0)
            
            # Check for error in metadata
            if "error" in metadata:
                result = {"error": metadata["error"]}
                result_queue.put(result)
                event_list[1].set()
                event_list[0].clear()
                continue
                
        except:
            metadata = {
                "batch_size": 1,
                "num_inputs": 100,
                "wall_time": 1.0,
                "init_ms": 0.0
            }
        
        # 通知推理进程数据已复制完成
        event_list[1].set()
        
        # 计算统计指标
        summary = compute_metrics(
            latencies=latencies,
            power_data=power_data,
            batch_size=metadata["batch_size"],
            num_inputs=metadata["num_inputs"],
            wall_time=metadata["wall_time"],
            init_ms=metadata["init_ms"]
        )
        
        # 将结果放入队列
        result_queue.put(summary)
        
        # 检查全局退出信号
        if event_list[2].is_set():
            break
        
        # 重置数据复制完成事件，准备下一轮
        event_list[1].clear()
        # 重置功耗数据就绪信号（如果存在）
        if len(event_list) > 3:
            event_list[3].clear()

